// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//

#include <nebula/compute/cast.h>

#include <mutex>
#include <sstream>
#include <string>
#include <unordered_map>
#include <unordered_set>
#include <utility>
#include <vector>

#include <nebula/compute/cast_internal.h>
#include <nebula/compute/exec.h>
#include <nebula/compute/function_internal.h>
#include <nebula/compute/kernel.h>
#include <nebula/compute/kernels/codegen_internal.h>
#include <nebula/compute/registry.h>
#include <turbo/log/logging.h>
#include <nebula/util/reflection_internal.h>

namespace nebula::compute {
    namespace internal {

        // ----------------------------------------------------------------------
        // Function options

        namespace {

            std::unordered_map<int, std::shared_ptr<CastFunction>> g_cast_table;
            std::once_flag cast_table_initialized;

            void AddCastFunctions(const std::vector<std::shared_ptr<CastFunction>> &funcs) {
                for (const auto &func: funcs) {
                    g_cast_table[static_cast<int>(func->out_type_id())] = func;
                }
            }

            void InitCastTable() {
                AddCastFunctions(GetBooleanCasts());
                AddCastFunctions(GetBinaryLikeCasts());
                AddCastFunctions(GetNestedCasts());
                AddCastFunctions(GetNumericCasts());
                AddCastFunctions(GetTemporalCasts());
                AddCastFunctions(GetDictionaryCasts());
                AddCastFunctions(GetExtensionCasts());
            }

            void EnsureInitCastTable() { std::call_once(cast_table_initialized, InitCastTable); }

            const FunctionDoc cast_doc{"Cast values to another data type",
                                       ("Behavior when values wouldn't fit in the target type\n"
                                        "can be controlled through CastOptions."),
                                       {"input"},
                                       "CastOptions"};

            // Metafunction for dispatching to appropriate CastFunction. This corresponds
            // to the standard SQL CAST(expr AS target_type)
            class CastMetaFunction : public MetaFunction {
            public:
                CastMetaFunction() : MetaFunction("cast", Arity::Unary(), cast_doc) {}

                turbo::Result<const CastOptions *> ValidateOptions(const FunctionOptions *options) const {
                    auto cast_options = static_cast<const CastOptions *>(options);

                    if (cast_options == nullptr || cast_options->to_type == nullptr) {
                        return turbo::invalid_argument_error(
                                "Cast requires that options be passed with "
                                "the to_type populated");
                    }

                    return cast_options;
                }

                turbo::Result<Datum> ExecuteImpl(const std::vector<Datum> &args,
                                                 const FunctionOptions *options,
                                                 ExecContext *ctx) const override {
                    TURBO_MOVE_OR_RAISE(auto cast_options, ValidateOptions(options));
                    // args[0].type() could be a nullptr so check for that before
                    // we do anything with it.
                    if (args[0].type() && args[0].type()->equals(*cast_options->to_type)) {
                        // Nested types might differ in field names but still be considered equal,
                        // so we can only return non-nested types as-is.
                        if (!is_nested(args[0].type()->id())) {
                            return args[0];
                        } else if (args[0].is_array()) {
                            TURBO_MOVE_OR_RAISE(std::shared_ptr<ArrayData> array,
                                                ::nebula::internal::get_array_view(
                                                        args[0].array(), cast_options->to_type.owned_type));
                            return Datum(array);
                        } else if (args[0].is_chunked_array()) {
                            TURBO_MOVE_OR_RAISE(
                                    std::shared_ptr<ChunkedArray> array,
                                    args[0].chunked_array()->View(cast_options->to_type.owned_type));
                            return Datum(array);
                        }
                    }

                    turbo::Result<std::shared_ptr<CastFunction>> result =
                            GetCastFunction(*cast_options->to_type);
                    if (!result.ok()) {
                        turbo::Status s = result.status();
                        return s.with_message(s.message(), " from ", *(args[0].type()));
                    }
                    return (*result)->execute(args, options, ctx);
                }
            };

            static auto kCastOptionsType = get_function_options_type<CastOptions>(
                    nebula::internal::DataMember("to_type", &CastOptions::to_type),
                    nebula::internal::DataMember("allow_int_overflow", &CastOptions::allow_int_overflow),
                    nebula::internal::DataMember("allow_time_truncate", &CastOptions::allow_time_truncate),
                    nebula::internal::DataMember("allow_time_overflow", &CastOptions::allow_time_overflow),
                    nebula::internal::DataMember("allow_decimal_truncate",
                                                 &CastOptions::allow_decimal_truncate),
                    nebula::internal::DataMember("allow_float_truncate",
                                                 &CastOptions::allow_float_truncate),
                    nebula::internal::DataMember("allow_invalid_utf8", &CastOptions::allow_invalid_utf8));
        }  // namespace

        void register_scalar_cast(FunctionRegistry *registry) {
                    KCHECK_OK(registry->add_function(std::make_shared<CastMetaFunction>()));
                    KCHECK_OK(registry->add_function_options_type(kCastOptionsType));
        }

        CastFunction::CastFunction(std::string name, Type::type out_type_id)
                : ScalarFunction(std::move(name), Arity::Unary(), FunctionDoc::Empty()),
                  out_type_id_(out_type_id) {}

        turbo::Status CastFunction::add_kernel(Type::type in_type_id, ScalarKernel kernel) {
            // We use the same KernelInit for every cast
            kernel.init = internal::CastState::init;
            TURBO_RETURN_NOT_OK(ScalarFunction::add_kernel(kernel));
            in_type_ids_.push_back(in_type_id);
            return turbo::OkStatus();
        }

        turbo::Status CastFunction::add_kernel(Type::type in_type_id, std::vector<InputType> in_types,
                                              OutputType out_type, ArrayKernelExec exec,
                                              NullHandling::type null_handling,
                                              MemAllocation::type mem_allocation) {
            ScalarKernel kernel;
            kernel.signature = KernelSignature::create(std::move(in_types), std::move(out_type));
            kernel.exec = exec;
            kernel.null_handling = null_handling;
            kernel.mem_allocation = mem_allocation;
            return add_kernel(in_type_id, std::move(kernel));
        }

        turbo::Result<const Kernel *> CastFunction::dispatch_exact(
                const std::vector<TypeHolder> &types) const {
            TURBO_RETURN_NOT_OK(check_arity(types.size()));

            std::vector<const ScalarKernel *> candidate_kernels;
            for (const auto &kernel: kernels_) {
                if (kernel.signature->matches_inputs(types)) {
                    candidate_kernels.push_back(&kernel);
                }
            }

            if (candidate_kernels.size() == 0) {
                return turbo::unimplemented_error("Unsupported cast from ", types[0].type->to_string(),
                                                  " to ", nebula::internal::to_type_name(out_type_id_), " using function ",
                                                  this->name());
            }

            if (candidate_kernels.size() == 1) {
                // One match, return it
                return candidate_kernels[0];
            }

            // Now we are in a casting scenario where we may have both a EXACT_TYPE and
            // a SAME_TYPE_ID. So we will see if there is an exact match among the
            // candidate kernels and if not we will just return the first one
            for (auto kernel: candidate_kernels) {
                const InputType &arg0 = kernel->signature->in_types()[0];
                if (arg0.kind() == InputType::EXACT_TYPE) {
                    // Bingo. Return it
                    return kernel;
                }
            }

            // We didn't find an exact match. So just return some kernel that matches
            return candidate_kernels[0];
        }

        turbo::Result<std::shared_ptr<CastFunction>> GetCastFunction(const DataType &to_type) {
            internal::EnsureInitCastTable();
            auto it = internal::g_cast_table.find(static_cast<int>(to_type.id()));
            if (it == internal::g_cast_table.end()) {
                return turbo::unimplemented_error("Unsupported cast to ", to_type);
            }
            return it->second;
        }

    }  // namespace internal

    CastOptions::CastOptions(bool safe)
            : FunctionOptions(internal::kCastOptionsType),
              allow_int_overflow(!safe),
              allow_time_truncate(!safe),
              allow_time_overflow(!safe),
              allow_decimal_truncate(!safe),
              allow_float_truncate(!safe),
              allow_invalid_utf8(!safe) {}

    bool CastOptions::is_safe() const {
        return !allow_int_overflow && !allow_time_truncate && !allow_time_overflow &&
               !allow_decimal_truncate && !allow_float_truncate && !allow_invalid_utf8;
    }

    bool CastOptions::is_unsafe() const {
        return allow_int_overflow && allow_time_truncate && allow_time_overflow &&
               allow_decimal_truncate && allow_float_truncate && allow_invalid_utf8;
    }

    constexpr char CastOptions::kTypeName[];

    turbo::Result<Datum> Cast(const Datum &value, const CastOptions &options, ExecContext *ctx) {
        return call_function("cast", {value}, &options, ctx);
    }

    turbo::Result<Datum> Cast(const Datum &value, const TypeHolder &to_type,
                              const CastOptions &options, ExecContext *ctx) {
        CastOptions options_with_to_type = options;
        options_with_to_type.to_type = to_type;
        return Cast(value, options_with_to_type, ctx);
    }

    turbo::Result<std::shared_ptr<Array>> Cast(const Array &value, const TypeHolder &to_type,
                                               const CastOptions &options, ExecContext *ctx) {
        TURBO_MOVE_OR_RAISE(Datum result, Cast(Datum(value), to_type, options, ctx));
        return result.make_array();
    }

    bool CanCast(const DataType &from_type, const DataType &to_type) {
        internal::EnsureInitCastTable();
        auto it = internal::g_cast_table.find(static_cast<int>(to_type.id()));
        if (it == internal::g_cast_table.end()) {
            return false;
        }

        const internal::CastFunction *function = it->second.get();
                DKCHECK_EQ(function->out_type_id(), to_type.id());

        for (auto from_id: function->in_type_ids()) {
            // XXX should probably check the output type as well
            if (from_type.id() == from_id) return true;
        }

        return false;
    }

}  // namespace nebula::compute
